import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from imblearn.over_sampling import SMOTE
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, VotingClassifier
from xgboost import XGBClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.svm import SVC
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score, roc_curve, auc
import os
from sklearn.feature_selection import SelectKBest, f_classif

# 全局绘图设置
sns.set(style="whitegrid", palette="pastel")
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def safe_drop(df, columns):
    return df.drop(columns=[col for col in columns if col in df.columns], errors='ignore')

def load_data():
    stroke = pd.read_csv('../preprogressData/output/stroke_processed.csv')
    heart = pd.read_csv('../preprogressData/output/heart_processed.csv')
    cirrhosis = pd.read_csv('../preprogressData/output/cirrhosis_processed.csv')
    return stroke, heart, cirrhosis

def preprocess_stroke(df):
    X = safe_drop(df, ['id', 'stroke'])
    y = df['stroke']
    X = pd.get_dummies(X)
    return X, y

def preprocess_heart(df):
    X = safe_drop(df, ['HeartDisease'])
    y = df['HeartDisease']
    X = pd.get_dummies(X)
    return X, y

def preprocess_cirrhosis(df, binary=True):
    X = safe_drop(df, ['ID', 'N_Days', 'Status', 'Stage'])
    y = df['Stage']
    if binary:
        y = (y >= 3).astype(int)
    X = pd.get_dummies(X)
    return X, y

def preprocess_cirrhosis_selectkbest(df, binary=True, k=10):
    X = safe_drop(df, ['ID', 'N_Days', 'Status', 'Stage'])
    y = df['Stage']
    if binary:
        y = (y >= 3).astype(int)
    X = pd.get_dummies(X)
    selector = SelectKBest(f_classif, k=min(k, X.shape[1]))
    X_new = selector.fit_transform(X, y)
    selected_features = X.columns[selector.get_support()]
    print('SelectKBest选中的特征:', list(selected_features))
    return pd.DataFrame(X_new, columns=selected_features), y

def plot_confusion_matrix(y_true, y_pred, classes, save_path, title):
    cm = confusion_matrix(y_true, y_pred)
    cm_sum = np.sum(cm)
    cm_perc = cm / cm_sum * 100 if cm_sum > 0 else cm
    annot = np.empty_like(cm).astype(str)
    nrows, ncols = cm.shape
    for i in range(nrows):
        for j in range(ncols):
            c = cm[i, j]
            p = cm_perc[i, j]
            if c == 0:
                annot[i, j] = ''
            else:
                annot[i, j] = f'{c}\n({p:.1f}%)'
    plt.figure(figsize=(6, 5))
    sns.heatmap(cm, annot=annot, fmt='', cmap='YlGnBu', xticklabels=classes, yticklabels=classes,
                annot_kws={'size': 16}, linewidths=1, linecolor='white', cbar=False, square=True)
    plt.xlabel('预测值', fontsize=14)
    plt.ylabel('真实值', fontsize=14)
    plt.title(title, fontsize=16, pad=15)
    plt.xticks(fontsize=13)
    plt.yticks(fontsize=13)
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.close()

def plot_roc_curve(y_true, y_score, save_path, title):
    plt.figure(figsize=(6, 5))
    fpr, tpr, _ = roc_curve(y_true, y_score)
    roc_auc = auc(fpr, tpr)
    plt.plot(fpr, tpr, color='#FF6F61', lw=3, label=f'ROC曲线 (AUC = {roc_auc:.3f})')
    plt.plot([0, 1], [0, 1], color='gray', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('假阳性率', fontsize=14)
    plt.ylabel('真阳性率', fontsize=14)
    plt.title(title, fontsize=16, pad=15)
    plt.legend(loc="lower right", fontsize=13)
    plt.grid(alpha=0.3)
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.close()

def train_and_evaluate_ensemble(X, y, disease_name, save_dir):
    os.makedirs(save_dir, exist_ok=True)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    # SMOTE重采样
    smote = SMOTE(random_state=42)
    res = smote.fit_resample(X_train_scaled, y_train)
    if isinstance(res, tuple) and len(res) == 2:
        X_train_res, y_train_res = res
    else:
        X_train_res, y_train_res = res[0], res[1]
    # 集成模型
    ensemble = VotingClassifier(estimators=[
        ('lr', LogisticRegression(max_iter=1000, random_state=42)),
        ('rf', RandomForestClassifier(random_state=42)),
        ('xgb', XGBClassifier(random_state=42, eval_metric='logloss', use_label_encoder=False)),
        ('dt', DecisionTreeClassifier(random_state=42)),
        ('svm', SVC(probability=True, random_state=42))
    ], voting='soft', n_jobs=-1)
    ensemble.fit(X_train_res, y_train_res)
    y_pred = ensemble.predict(X_test_scaled)
    y_score = ensemble.predict_proba(X_test_scaled)[:, 1]
    # 评估
    report = classification_report(y_test, y_pred, output_dict=True)
    pd.DataFrame(report).to_csv(os.path.join(save_dir, f'{disease_name}_ensemble_report.csv'))
    class_names = ['健康', '患病']
    plot_confusion_matrix(y_test, y_pred, class_names, os.path.join(save_dir, f'{disease_name}_ensemble_cm.png'), f'{disease_name} 集成模型混淆矩阵')
    auc_score = roc_auc_score(y_test, y_score)
    plot_roc_curve(y_test, y_score, os.path.join(save_dir, f'{disease_name}_ensemble_roc.png'), f'{disease_name} 集成模型ROC曲线 (AUC={auc_score:.2f})')
    print(f"\n{disease_name} 集成模型AUC: {auc_score:.4f}")
    print(classification_report(y_test, y_pred, digits=4))
    return auc_score

def main():
    os.makedirs('new_results/stroke', exist_ok=True)
    os.makedirs('new_results/heart', exist_ok=True)
    os.makedirs('new_results/cirrhosis', exist_ok=True)
    stroke, heart, cirrhosis = load_data()
    print("\n================ 集成模型：中风 ================")
    X_s, y_s = preprocess_stroke(stroke)
    train_and_evaluate_ensemble(X_s, y_s, 'stroke', 'new_results/stroke')
    print("\n================ 集成模型：心脏病 ================")
    X_h, y_h = preprocess_heart(heart)
    train_and_evaluate_ensemble(X_h, y_h, 'heart', 'new_results/heart')
    print("\n================ 集成模型：肝硬化（全部特征） ================")
    X_c, y_c = preprocess_cirrhosis(cirrhosis, binary=True)
    train_and_evaluate_ensemble(X_c, y_c, 'cirrhosis_all', 'new_results/cirrhosis')
    print("\n================ 集成模型：肝硬化（SelectKBest特征） ================")
    X_c_k, y_c_k = preprocess_cirrhosis_selectkbest(cirrhosis, binary=True, k=10)
    train_and_evaluate_ensemble(X_c_k, y_c_k, 'cirrhosis_kbest', 'new_results/cirrhosis')

if __name__ == '__main__':
    main()
